# link: https://mp.weixin.qq.com/s/NTct2_AYhz4Z8q5MYtBQcA

import backtrader as bt
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# 贵州茅台
data2 = get_data(codes[1])
# 海天味业
data3 = get_data(codes[2])

# %%

# 实例化策略
cerebro = bt.Cerebro()

st_date = datetime.datetime(2023, 4, 4)
ed_date = datetime.datetime(2023, 5, 11)

# 添加 600276.SH 的行情数据
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name=codes[0])

# 添加 600519.SH 的行情数据
datafeed2 = bt.feeds.PandasData(dataname=data2, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed2, name=codes[1])

# 添加 603288.SH 的行情数据
datafeed3 = bt.feeds.PandasData(dataname=data3, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed3, name=codes[2])


# 第3节：提取 line 上的数据点，使用 get(ago,size) 切片函数
class TestStrategy(bt.Strategy):
    def __init__(self):
        self.count = 0  # 用于计算 next 的循环次数
        # 打印数据集和数据集对应的名称
        print("------------- init 中的索引位置-------------")
        # 对 datetime 线进行索引时，xxx.date(X) 可以直接以“xxxx-xx-xx xx:xx:xx”的形式返回，X 就是索引位置，可以看做是传统 [X] 索引方式的改进版
        print("0 索引：", 'datetime', self.data1.lines.datetime.date(0), 'close', self.data1.lines.close[0])
        print("-1 索引：", 'datetime', self.data1.lines.datetime.date(-1), 'close', self.data1.lines.close[-1])
        print("-2 索引", 'datetime', self.data1.lines.datetime.date(-2), 'close', self.data1.lines.close[-2])
        print("1 索引：", 'datetime', self.data1.lines.datetime.date(1), 'close', self.data1.lines.close[1])
        print("2 索引", 'datetime', self.data1.lines.datetime.date(2), 'close', self.data1.lines.close[2])
        # 通过 get() 切片时，如果是从 ago=0 开始取，不会返回数据，从其他索引位置开始取，能返回数据
        print("从 0 开始往前取3天的收盘价：", self.data1.lines.close.get(ago=0, size=3))
        print("从-1开始往前取3天的收盘价：", self.data1.lines.close.get(ago=-1, size=3))
        print("从-2开始往前取3天的收盘价：", self.data1.lines.close.get(ago=-2, size=3))
        print("line的总长度：", self.data1.buflen())

    def next(self):
        print(f"------------- next 的第{self.count + 1}次循环 --------------")
        print("当前时点（今日）：", 'datetime', self.data1.lines.datetime.date(0), 'close', self.data1.lines.close[0])
        print("往前推1天（昨日）：", 'datetime', self.data1.lines.datetime.date(-1), 'close', self.data1.lines.close[-1])
        print("往前推2天（前日）", 'datetime', self.data1.lines.datetime.date(-2), 'close', self.data1.lines.close[-2])
        print("前日、昨日、今日的收盘价：", self.data1.lines.close.get(ago=0, size=3))
        print("往后推1天（明日）：", 'datetime', self.data1.lines.datetime.date(1), 'close', self.data1.lines.close[1])
        print("往后推2天（明后日）", 'datetime', self.data1.lines.datetime.date(2), 'close', self.data1.lines.close[2])
        # 在 next() 中调用 len(self.data0)，返回的是当前已处理（已回测）的数据长度，会随着回测的推进动态增长
        print("已处理的数据点：", len(self.data1))
        # buflen() 返回整条线的总长度，固定不变；
        print("line的总长度：", self.data0.buflen())
        self.count += 1


cerebro.addstrategy(TestStrategy)
result = cerebro.run()
